package trace

import (
	"context"
	"strings"
	"sync"
	"time"

	"xp/lib/errors"
	"xp/lib/etc"
	"xp/lib/jsonx"
	"xp/lib/md"
	"xp/lib/utils"
)

const (
	XReqID = `X-ND-Req-ID`
	XTrace = `X-ND-Trace`
)

var (
	traceEnable = etc.Bool("trace", "enable")
)

func EnableTrace(enable bool) {
	traceEnable = enable
}

type Span struct {
	TraceID  string     `json:"-"`
	ID       int64      `json:"id,string"`
	Name     string     `json:"name"`
	Start    jsonx.Time `json:"start"`
	End      jsonx.Time `json:"end"`
	Labels   []string   `json:"labels" jsonx:"unsafe"`
	ParentID int64      `json:"parent_id,string"`
	Error    []byte     `json:"error"`

	wg *sync.WaitGroup
}

func NewSpan(ctx context.Context, name string, labels ...string) *Span {
	if ctx == nil {
		return nil
	}
	trace, span := FromContext(ctx)
	if trace == nil || span == nil {
		return nil
	}
	ctxLables := Labels(ctx)
	if len(ctxLables) > 0 {
		labels = append(labels, "ctxLables", strings.Join(ctxLables, ","))
	}
	return trace.NewSpan(span.ID, name, labels...)
}

type spanCtxKey struct{}

func (s *Span) NewContext(ctx context.Context) context.Context {
	return context.WithValue(ctx, spanCtxKey{}, s)
}

func (s *Span) Finish(err error) {
	s.End = jsonx.Time(time.Now())
	s.Error = errors.ToBytes(err)
	s.wg.Done()
}

func (s *Span) SetKeyvals(kvs ...string) {
	s.Labels = append(s.Labels, kvs...)
}

type Trace struct {
	*sync.Mutex

	ID       string  `json:"id"`
	Host     string  `json:"host"`
	Port     int     `json:"port,string"`
	Service  string  `json:"service"`
	Revision string  `json:"revision"`
	Spans    []*Span `json:"spans"`

	wg *sync.WaitGroup
}

func New(parent int64, id, name string, labels ...string) *Trace {
	if !traceEnable {
		return nil
	}
	tr := &Trace{
		ID:       id,
		Mutex:    new(sync.Mutex),
		Host:     md.GetHost(),
		Port:     md.Get().Host.Port,
		Service:  md.GetService(),
		Revision: md.Get().VersionControl.Revision,
		wg:       new(sync.WaitGroup),
	}
	tr.NewSpan(parent, name, labels...)
	return tr
}

type traceCtxkey struct{}

func (t *Trace) NewContext(ctx context.Context) context.Context {
	ctx = context.WithValue(ctx, traceCtxkey{}, t)
	return t.Spans[0].NewContext(ctx)
}

func (t *Trace) Finish(err error) {
	t.Spans[0].Finish(err)
	t.wg.Wait()
	var buf = make([]byte, t.Len())
	n, _ := t.MarshalJsonx(buf)
	writer.Write(buf[:n])
}

func (t *Trace) NewSpan(parent int64, name string, labels ...string) *Span {
	span := &Span{
		TraceID:  t.ID,
		Name:     name,
		Start:    jsonx.Time(time.Now()),
		ParentID: parent,
		Labels:   labels,
		wg:       t.wg,
	}
	t.wg.Add(1)
	t.Lock()
	span.ID = utils.NewGUID()
	t.Spans = append(t.Spans, span)
	t.Unlock()
	return span
}

func (t *Trace) SetKeyvals(kvs ...string) {
	t.Spans[0].SetKeyvals(kvs...)
}

func FromContext(ctx context.Context) (trace *Trace, parent *Span) {
	tr := ctx.Value(traceCtxkey{})
	if tr != nil {
		trace = tr.(*Trace)
	}
	s := ctx.Value(spanCtxKey{})
	if s != nil {
		parent = s.(*Span)
	}
	return
}

type labelsCtxKey struct{}

func Labels(ctx context.Context) []string {
	origin := ctx.Value(labelsCtxKey{})
	if origin == nil {
		return []string{}
	}
	return origin.([]string)
}

func WithLables(ctx context.Context, labels ...string) context.Context {
	return context.WithValue(ctx, labelsCtxKey{}, append(Labels(ctx), labels...))
}
